package com.akto.action.gpt.handlers;

import com.akto.action.gpt.GptAction;
import com.akto.action.gpt.result_fetchers.ResultFetcherStrategy;
import com.akto.log.LoggerMaker;
import com.mongodb.BasicDBObject;

public class AnalyzeVulnerability implements QueryHandler {
    
    private static final LoggerMaker loggerMaker = new LoggerMaker(AnalyzeVulnerability.class, LoggerMaker.LogDb.DASHBOARD);
    private final ResultFetcherStrategy<BasicDBObject> resultFetcherStrategy;
    
    public AnalyzeVulnerability(ResultFetcherStrategy<BasicDBObject> resultFetcherStrategy) {
        this.resultFetcherStrategy = resultFetcherStrategy;
    }
    
    @Override
    public BasicDBObject handleQuery(BasicDBObject meta) throws Exception {
        BasicDBObject testData = (BasicDBObject) meta.get("contentData");
        
        // If testData is not provided, this is an error
        if (testData == null) {
            throw new Exception("No test data provided for vulnerability analysis");
        }
        
        
        BasicDBObject request = new BasicDBObject();
        request.put("query_type", GptQuery.ANALYZE_VULNERABILITY.getName());
        request.put("test_data", testData);  // Pass structured data instead of raw response
        request.put(GptAction.USER_EMAIL, meta.getString(GptAction.USER_EMAIL));
        
        // Create the prompt for LLM using the structured data
        String prompt = buildPrompt(testData);
        request.put("prompt", prompt);
        
        // Call LLM to analyze the response
        BasicDBObject llmResponse = this.resultFetcherStrategy.fetchResult(request);
        
        // Log the prompt being sent to LLM
        loggerMaker.infoAndAddToDb("Sending prompt to LLM for vulnerability analysis: " + prompt, LoggerMaker.LogDb.DASHBOARD);
        // Log the LLM response for debugging/monitoring
        if (llmResponse != null) {
            loggerMaker.infoAndAddToDb("LLM Response received for vulnerability analysis: " + llmResponse, LoggerMaker.LogDb.DASHBOARD);
        } else {
            loggerMaker.infoAndAddToDb("No LLM response received for vulnerability analysis", LoggerMaker.LogDb.DASHBOARD);
        }
        
        // The resultFetcherStrategy already returns the processed vulnerability analysis
        // No need for additional processing - return the result directly
        return llmResponse != null ? llmResponse : new BasicDBObject();
        //return createDummyVulnerabilityResult();
    }
    
    private String buildPrompt(BasicDBObject testData) {
        StringBuilder prompt = new StringBuilder();
        
        // Extract structured data
        BasicDBObject testContext = (BasicDBObject) testData.get("testContext");
        BasicDBObject request = (BasicDBObject) testData.get("request");
        BasicDBObject response = (BasicDBObject) testData.get("response");
        
        // Get test category from the structured data
        String testCategory = "UNKNOWN";
        if (testContext != null && testContext.getString("category") != null) {
            testCategory = testContext.getString("category");
        }
        
        // Use testCategory as the primary vulnerability type
        prompt.append("Analyze the following HTTP request/response for ").append(testCategory).append(" vulnerabilities.\n\n");
        
        
        // TODO: Optimize repetitive if-conditions by using a map of fields to extract

        // Add test context information
        if (testContext != null) {
            prompt.append("Test Context:\n");
            if (testContext.getString("category") != null) {
                prompt.append("- Category: ").append(testContext.getString("category")).append("\n");
            }
            if (testContext.getString("description") != null) {
                prompt.append("- Description: ").append(testContext.getString("description")).append("\n");
            }
            if (testContext.getString("severity") != null) {
                prompt.append("- Severity: ").append(testContext.getString("severity")).append("\n");
            }
            if (testContext.get("cwe") != null) {
                prompt.append("- CWE IDs: ").append(testContext.get("cwe")).append("\n");
            }
            prompt.append("\n");
        }
        
        // Add minimal request context
        if (request != null) {
            prompt.append("Request Context:\n");
            if (request.getString("method") != null) {
                prompt.append("- Method: ").append(request.getString("method")).append("\n");
            }
            if (request.getString("url") != null) {
                prompt.append("- URL: ").append(request.getString("url")).append("\n");
            }
            if (request.getString("contentType") != null) {
                prompt.append("- Content-Type: ").append(request.getString("contentType")).append("\n");
            }
            if (request.getString("authorization") != null) {
                prompt.append("- Authorization: ").append(request.getString("authorization")).append("\n");
            }
            if (request.getString("userAgent") != null) {
                prompt.append("- User-Agent: ").append(request.getString("userAgent")).append("\n");
            }
            if (request.getString("xForwardedFor") != null) {
                prompt.append("- X-Forwarded-For: ").append(request.getString("xForwardedFor")).append("\n");
            }
            if (request.getString("origin") != null) {
                prompt.append("- Origin: ").append(request.getString("origin")).append("\n");
            }
            if (request.getString("referer") != null) {
                prompt.append("- Referer: ").append(request.getString("referer")).append("\n");
            }
            prompt.append("\n");
        }
        
        // Add response data (primary analysis target) - formatted exactly like frontend editor
        if (response != null) {
            prompt.append("Response Data to analyze:\n");
            prompt.append("```\n");
            
            // Format as HTTP response (exactly same as frontend editor)
            if (response.getInt("statusCode") > 0) {
                prompt.append(response.getInt("statusCode")).append("\n\n");
            }
            
            if (response.get("headers") != null) {
                prompt.append(response.get("headers")).append("\n\n");
            }
            
            if (response.getString("body") != null) {
                prompt.append(response.getString("body"));
            }
            prompt.append("\n```\n\n");
        }
        
        prompt.append("Instructions:\n");
        prompt.append("1. Focus on the response data for ").append(testCategory).append(" vulnerability indicators.\n");
        prompt.append("   - If you find explicit indicators (e.g., payload reflected, JWT 'alg':'none', SQL errors, prompt injection instructions, etc.), identify the exact continuous substring(s) that are vulnerable.\n");
        prompt.append("   - If the vulnerability manifests as a lack of enforcement (e.g., response succeeds when it should fail, authentication or input validation is missing), mark the entire response data as vulnerable.\n");
        prompt.append("2. CRITICAL: Character positions (start/end) must be calculated from the ENTIRE response data above, including status code, headers, and body. The response data includes status code, headers, and body in the exact format shown.\n");
        prompt.append("3. For each vulnerability found, provide the start and end position in the complete response data above (character index starting from 0).\n");
        prompt.append("4. Consider the request context (method, URL, content-type) when analyzing the response.\n");
        prompt.append("5. IMPORTANT: The phrase must be a continuous substring from the response data. If a vulnerability spans multiple non-continuous parts, return multiple segments instead of one broken segment.\n");
        prompt.append("6. Return the analysis in JSON format with structure:\n");
        
        prompt.append("{\n");
        prompt.append("  \"vulnerableSegments\": [\n");
        prompt.append("    {\n");
        prompt.append("      \"start\": number,\n");
        prompt.append("      \"end\": number,\n");
        prompt.append("      \"phrase\": \"exact continuous vulnerable substring\"\n");
        prompt.append("    }\n");
        prompt.append("  ]\n");
        prompt.append("}\n\n");
        
        return prompt.toString();
    }
    

    /**
     * Creates dummy vulnerability analysis result for testing
     */
    private BasicDBObject createDummyVulnerabilityResult() {
        BasicDBObject result = new BasicDBObject();
        
        // Create dummy vulnerable segments
        BasicDBObject[] vulnerableSegments = {
            new BasicDBObject()
                .append("start", 15)
                .append("end", 45),
            new BasicDBObject()
                .append("start", 60)
                .append("end", 85),
            new BasicDBObject()
                .append("start", 120)
                .append("end", 150)
        };
        
        result.put("vulnerableSegments", vulnerableSegments);
        result.put("isDummyData", true);
        result.put("timestamp", System.currentTimeMillis());
        
        return result;
    }
}